"""
@Date: 2022/3/16
@Author: pyn
"""
import sqlite3, random
from typing import Tuple
import numpy as np
import pandas as pd
from gym.spaces import Box
from wtpy import WTSBarStruct
from wtpy.ExtModuleDefs import BaseExtDataLoader
from wtpy.StrategyDefs import BaseCtaStrategy, CtaContext, HftContext
from wenv.envs import WtpyCtaEnv, ActionPlan, RewardPlan, Stopper, Observer, Informer
from wenv.envs import EngineType, Space, Discrete


def sqlite2df(db_addr, table_name):
    conn = sqlite3.connect(db_addr)
    return pd.read_sql("select * from %s" % (table_name), conn)


class MyDataloader_d(BaseExtDataLoader):
    def load_final_his_bars(self, stdCode: str, period: str, feeder) -> bool:
        """"""
        print("enter MyDataloader...")
        params2tbname = {
            'IFd1': 'ifd', 'IFm1': 'if1m', 'SRd1': 'srd', 'SRm1': 'sr1m'
        }
        df = sqlite2df('./config/srif.db', params2tbname.get(stdCode.split('.')[1] + period))

        def process_and_feeder(df, **kwargs):
            df['Unnamed: 0'] = df['Unnamed: 0'].str[0:19]
            df['date'] = df.apply(lambda x: x['Unnamed: 0'][0:10], axis=1)
            df['date'] = df['date'].str.replace('-', '/')
            if 'time_value' in kwargs:
                df['time'] = kwargs['time_value']
            else:
                df['time'] = df.apply(lambda x: x['Unnamed: 0'][-8:], axis=1)
            columns_convert = {
                'date': 'date', 'time': 'time', 'open': 'open', 'high': 'high', 'low': 'low', 'close': 'close',
                'volume': 'vol'
            }
            df = df.rename(columns=columns_convert)
            df['date'] = df['date'].astype('datetime64').dt.strftime('%Y%m%d').astype('int64')
            df['time'] = (df['date'] - 19900000) * 10000 + df['time'].str.replace(':', '').str[:-2].astype('int')
            BUFFER = WTSBarStruct * len(df)
            buffer = BUFFER()

            def assign(procession, buffer):
                tuple(map(lambda x: setattr(buffer[x[0]], procession.name, x[1]), enumerate(procession)))

            df.apply(assign, buffer=buffer)
            feeder(buffer, len(df))

        if period == 'd1':
            process_and_feeder(df, time_value='00:00:00')
        elif period == 'm1':
            process_and_feeder(df)

        return True


class DefaultStra(BaseCtaStrategy, ActionPlan):
    """"""

    def __init__(self, name):
        """"""
        super().__init__(name=name)

    def engine_type(self) -> EngineType:
        return EngineType.ET_CTA

    def save_action(self, action):
        """"""
        self._action_ = action

    @property
    def action_space(self) -> Space:
        """"""
        return Discrete(3)

    def on_init(self, context: CtaContext):
        """"""
        context.stra_get_bars(stdCode='CZCE.SR.HOT', period='d1', count=50, isMain=True)
        self._symbols_ = ['CZCE.SR.HOT']

    def on_calculate(self, context: CtaContext):
        """"""
        print('on calc ...')
        self._env._observer.get_obs(context=context)
        self._env._rewardplan.get_reward(context=context)

    def on_calculate_done(self, context: CtaContext):
        """"""
        print('on calc done ...')
        action2position = {
            0: 1, 1: 0, 2: -1
        }
        for code in self._symbols_:
            context.stra_set_position(stdCode=code, qty=action2position[self._action_])
        # TODO action的映射


class DefaultRewardPlan(RewardPlan):
    """"""

    def __init__(self):
        """"""
        self._his_rewards = []

    def reset(self):
        """"""
        self._his_rewards = []

    def get_reward(self, context: CtaContext):
        dynbalance = context.stra_get_fund_data(0)  # 动态权益
        # closeprofitloss = context.stra_get_fund_data(1)  # 总平仓盈亏
        # positionprofitloss = context.stra_get_fund_data(2)  # 总浮动盈亏
        # totalfee = context.stra_get_fund_data(3)  # 总手续费
        self._his_rewards.append(dynbalance)

    @property
    def reward(self) -> float:
        """"""
        return self._his_rewards[-1]


class DefaultObserver(Observer):
    """"""

    def __init__(self):
        self.load_final_his_observations()

    def load_final_his_observations(self):
        """"""
        df = pd.read_excel('./rawdata/srd.xlsx')
        df['Unnamed: 0'] = df['Unnamed: 0'].str[0:19]
        df['date'] = df.apply(lambda x: x['Unnamed: 0'][0:10], axis=1)
        df['date'] = df['date'].str.replace('-', '/')

        df['time'] = df.apply(lambda x: x['Unnamed: 0'][-8:], axis=1)
        df['date'] = df['date'].astype('datetime64').dt.strftime('%Y%m%d').astype('int64')
        df['time'] = (df['date'] - 19900000) * 10000 + df['time'].str.replace(':', '').str[:-2].astype('int')
        df['bartimes'] = df['date'] * 10000 + 1500
        df = df[['bartimes', 'money', 'volume']]

        def list_split(listTemp, n):
            for i in range(0, len(listTemp), n):
                yield listTemp[i:i + n]

        self._timeranges_ = list(list_split(list(df['bartimes']), 10))
        self._observations_ = df.set_index('bartimes').T.to_dict('list')

    def get_obs(self, context: CtaContext):
        self._internal_obs_ = np.array([context.stra_get_fund_data(0), context.stra_get_fund_data(3)])
        self._time_ = context.stra_get_date() * 10000 + context.stra_get_time()

    @property
    def observation_space(self) -> Space:
        """"""
        return Box(**dict(low=-np.inf, high=np.inf, shape=(4,), dtype=np.float64))

    @property
    def obs(self) -> np.ndarray:
        """"""
        external = np.array(self._observations_.get(self._time_, []))
        return np.append(external, self._internal_obs_)

    def pick_timerange(self) -> Tuple[int, int]:
        return random.sample(self._timeranges_, 1)[0]


class DefaultStopper(Stopper):
    """"""

    @property
    def is_stop(self) -> bool:
        # TODO 传context?
        return False  # 完全由timerange控制episode结束


class DefaultInformer(Informer):
    """"""

    @property
    def info(self) -> dict:
        # TODO 传context?
        return {}


if __name__ == '__main__':
    observer = DefaultObserver()
    env = WtpyCtaEnv(
        extloader=MyDataloader_d(),
        Strategy=DefaultStra,
        rewardplan=DefaultRewardPlan(),
        observer=observer,
        stopper=DefaultStopper(),
        informer=DefaultInformer()
    )

    firstobs = env.reset()
    print(firstobs, type(firstobs))
    for i in range(10):
        action = env.action_space.sample()
        obs, reward, done, info = env.step(action)
        print('step:', i, obs, reward, done, info)
